#! /usr/bin/env python3
import paramiko
from paramiko.ssh_exception import NoValidConnectionsError
from paramiko.ssh_exception import AuthenticationException
from paramiko.ssh_exception import SSHException
import os

import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))

from utils.custom_log import log_start
logger = log_start('sshops')

class SSHConnection():
    """
    SSH连接类
    """
    def __init__(self,conndict,mode='key'):
        self.mode = mode
        self.host = conndict["hostip"]
        self.port = conndict["sshport"]
        self.nkey = conndict["sshkeyfile"]
        self.username = conndict["sshuser"]
        self.sshpwd = conndict.get("sshpwd", "")
        self.timeout = 60


    def open(self):
        
        info = ">> {2}@{1}  port ['{0}']".format(self.port,self.host,self.username)
        logger.info(info)

        if self.mode == "key":
            self.connect_key()
        else:
            self.connect_pwd()


    def close(self):
        try:
            self.sshclient.close()
        except Exception as e:
            error_msg = "关闭SSH连接失败：{}，请检查！".format(str(e))
            logger.error(error_msg)
            raise ValueError(error_msg)




    def connect_key(self):
        try:
            self.private_key = paramiko.RSAKey.from_private_key_file(self.nkey)
            self.sshclient = paramiko.SSHClient()
            self.sshclient.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            self.sshclient.connect(hostname=self.host, port=self.port, username=self.username, pkey=self.private_key, allow_agent=False, look_for_keys=False)
            return self.sshclient
        except SSHException as ssh:
            error_msg = "SSH连接失败：{}，请检查重试！".format(str(ssh))
            logger.error(error_msg)
            raise ValueError(error_msg)

    def connect_pwd(self):
        try:

            self.sshclient = paramiko.SSHClient()
            self.sshclient.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            self.sshclient.connect(hostname=self.host, port=self.port, username=self.username, password=self.sshpwd, allow_agent=False, look_for_keys=False)
            return self.sshclient
        except SSHException as ssh:
            error_msg = "SSH连接失败：{}，请检查！".format(str(ssh))
            logger.error(error_msg)
            raise ValueError(error_msg)

    
    # def _get_ttytype(self):
    #     tty_type = False
    #     if self.username != 'root':
    #         check_cmd = "sudo cat /etc/sudoers|grep -E 'requiretty' | grep -v 'See requiretty above'"
    #         stdin, stdout, stderr = self.sshclient.exec_command(check_cmd, get_pty=True)
    #         strinfo = stdout.read().decode().strip()
    #         if strinfo != '':
    #             if "#" not in strinfo:
    #                 tty_type = True
    #     return tty_type


    def exec_cmd(self, cmd,show=True):
        # relttytype = self._get_ttytype()
        # stdin, stdout, stderr = self.sshclient.exec_command(cmd, get_pty=relttytype)
        stdin, stdout, stderr = self.sshclient.exec_command(cmd)
        cmd_result = stdout.read().decode()
        channel = stdout.channel
        status = channel.recv_exit_status()
        if cmd_result or status == 0:
            if show:
                logger.info("远程执行命令：%s" % cmd)
            return  cmd_result.strip()
        else:
            error_info = stderr.read().decode().strip()
            error_msg = "远程执行失败：%s；=======>返回：%s" % (cmd,error_info)
            logger.error(error_msg)
            raise ValueError(error_msg)

    def exec_cmd_sync(self, cmd):
        try:
            def line_buffered(f):
                line_buf = ""
                while not f.channel.exit_status_ready():
                    line_buf += f.read(1).decode(encoding="utf-8",errors="ignore")
                    if line_buf.endswith('\n'):
                        yield line_buf
                        line_buf = ''
            stdin, stdout, stderr = self.sshclient.exec_command(cmd)
            for l in line_buffered(stdout):
                logger.info(l)

        except Exception as e:
            error_msg = "远程执行失败：%s；=======>返回：%s" % (cmd,e)
            logger.error(error_msg)
            raise ValueError(error_msg)


    def sftp_put_file(self, local_file, remote_dir):
        try:
            transport = self.sshclient.get_transport()
            sftp = paramiko.SFTPClient.from_transport(transport)
            remote_file = os.path.join(remote_dir, os.path.basename(local_file))
            sftp.put(local_file, remote_file)

            logger.info('上传本地文件命令：%s'% local_file)
        except Exception as e:
            error_msg = '上传本地文件失败：%s; =======>返回：%s'% (local_file,str(e))
            logger.error(error_msg)
            raise ValueError(error_msg)



    def sftp_get_file(self, file, local_dir, remote_file):
        try:
            transport = self.sshclient.get_transport()
            sftp = paramiko.SFTPClient.from_transport(transport)
            sftp.get(remote_file, os.path.join(local_dir, file))
            logger.info('下载远端文件命令：%s' % remote_file)
        except Exception as e:
            error_msg = '下载远端文件失败：%s'% remote_file
            logger.error(error_msg)
            raise ValueError(error_msg)


    def __enter__(self):
        self.open()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

if __name__ == '__main__':
    conndict = {
        "hostip": "192.168.0.110",
        "sshport": "22",
        "sshuser": "root",
        "sshpwd": "123456",
        'sshkeyfile': '/Users/lijingjing/.ssh/id_rsa'
    }

    with SSHConnection(conndict,mode='key') as ssh:
        v = ssh.exec_cmd('date',show=False)
        logger.info(v)